import os
import re
import winreg
import zipfile
import requests
import subprocess
from io import BytesIO
from packaging import version


def get_chrome_version():
    """
    获取 Windows 系统上安装的 Google Chrome 版本
    """
    try:
        # 方法1：通过注册表获取 Chrome 版本（推荐）
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Google\Chrome\BLBeacon") as key:
            chrome_version, _ = winreg.QueryValueEx(key, "version")
            return chrome_version
    except:
        pass

    try:
        # 方法2：通过查询 Chrome 可执行文件路径获取版本
        chrome_path = r"C:\Program Files\Google\Chrome\Application\chrome.exe"
        if not os.path.exists(chrome_path):
            chrome_path = r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe"

        if os.path.exists(chrome_path):
            # 使用 PowerShell 获取文件版本信息
            ps_command = f'(Get-Item "{chrome_path}").VersionInfo.ProductVersion'
            result = subprocess.run(["powershell", "-Command", ps_command], capture_output=True, text=True)
            return result.stdout.strip()
    except:
        pass

    try:
        # 方法3：尝试执行 chrome --version 命令
        result = subprocess.run(["chrome", "--version"], capture_output=True, text=True)
        if result.returncode == 0:
            match = re.search(r"\d+\.\d+\.\d+\.\d+", result.stdout)
            if match:
                return match.group(0)
    except:
        pass

    raise Exception("无法检测到 Chrome 安装或版本信息")


def get_major_version(full_version):
    """
    从完整版本号中提取主版本号
    """
    # 使用正则表达式提取主版本号
    match = re.match(r"^(\d+)\.\d+\.\d+\.\d+", full_version)
    if match:
        return int(match.group(1))

    # 如果正则失败，尝试使用版本解析库
    try:
        return int(str(version.parse(full_version).major))
    except:
        raise ValueError(f"无法解析版本号: {full_version}")


def get_chromedriver_version(chrome_major_version):
    """
    获取与 Chrome 主版本匹配的 ChromeDriver 版本
    """
    # ChromeDriver 版本检查 API
    url = f"https://mirrors.huaweicloud.com/chromedriver/{chrome_major_version}/"

    try:
        response = requests.get(url)
        response.raise_for_status()
        return response.text.strip()
    except requests.exceptions.RequestException:
        # 如果指定主版本不存在，尝试获取最新稳定版本
        latest_url = "https://chromedriver.storage.googleapis.com/LATEST_RELEASE"
        response = requests.get(latest_url)
        response.raise_for_status()
        latest_version = response.text.strip()

        # 检查最新版本是否兼容
        latest_major = get_major_version(latest_version)
        if latest_major == chrome_major_version:
            return latest_version

        # 如果不兼容，尝试获取兼容版本
        compatible_url = f"https://chromedriver.storage.googleapis.com/LATEST_RELEASE_{chrome_major_version}"
        response = requests.get(compatible_url)
        if response.status_code == 200:
            return response.text.strip()

        raise Exception(f"找不到与 Chrome 主版本 {chrome_major_version} 兼容的 ChromeDriver")


def download_chromedriver(chromedriver_version, save_path="."):
    """
    下载并解压 ChromeDriver
    """
    # 创建下载目录
    os.makedirs(save_path, exist_ok=True)
    ver_arr = chromedriver_version.split(".")
    pro_three_version = ver_arr[0] + '.' + ver_arr[1] + '.' + ver_arr[2] + '.%d'
    print(ver_arr[3])
    print(type(ver_arr[3]))
    last_num = int(ver_arr[3])
    for i in range(last_num):
        chromedriver_version = pro_three_version % i
        print(chromedriver_version)
        # 下载 URL
        download_url = f"https://mirrors.huaweicloud.com/chromedriver/{chromedriver_version}/chromedriver-win64.zip"

        try:
            print(f"正在下载 ChromeDriver {chromedriver_version}...")
            response = requests.get(download_url)
            response.raise_for_status()

            # 解压 ZIP 文件
            with zipfile.ZipFile(BytesIO(response.content)) as zip_file:
                # 提取 chromedriver.exe
                zip_file.extract("/chromedriver-win64/chromedriver.exe", save_path)
                chromedriver_path = os.path.join(save_path, "chromedriver.exe")

                # 验证文件是否成功提取
                if os.path.exists(chromedriver_path):
                    print(f"成功下载 ChromeDriver 到: {chromedriver_path}")
                    return chromedriver_path
                else:
                    raise Exception("解压后未找到 chromedriver.exe")
        except Exception as e:
            raise Exception(f"下载 ChromeDriver 失败: {str(e)}")


def verify_chromedriver(chrome_version, chromedriver_path):
    """
    验证 ChromeDriver 与 Chrome 版本是否匹配
    """
    try:
        # 获取 ChromeDriver 版本
        result = subprocess.run([chromedriver_path, "--version"], capture_output=True, text=True)
        if result.returncode != 0:
            return False, f"无法获取 ChromeDriver 版本: {result.stderr}"

        # 提取版本号
        match = re.search(r"(\d+\.\d+\.\d+\.\d+)", result.stdout)
        if not match:
            return False, "无法解析 ChromeDriver 版本"

        chromedriver_version = match.group(1)

        # 比较主版本号
        chrome_major = get_major_version(chrome_version)
        driver_major = get_major_version(chromedriver_version)

        if chrome_major == driver_major:
            return True, f"版本验证成功: Chrome {chrome_version} 匹配 ChromeDriver {chromedriver_version}"
        else:
            return False, f"版本不匹配: Chrome {chrome_version} (主版本 {chrome_major}) vs ChromeDriver {chromedriver_version} (主版本 {driver_major})"
    except Exception as e:
        return False, f"验证失败: {str(e)}"


def main():
    try:
        # 1. 获取 Chrome 版本
        chrome_version = get_chrome_version()
        print(f"检测到 Chrome 版本: {chrome_version}")

        # 2. 提取主版本号
        chrome_major = get_major_version(chrome_version)
        print(f"Chrome 主版本: {chrome_major}")

        # 3. 获取匹配的 ChromeDriver 版本
        # chromedriver_version = get_chromedriver_version(chrome_version)
        # print(f"匹配的 ChromeDriver 版本: {chromedriver_version}")

        # 4. 下载 ChromeDriver
        download_dir = os.path.join(os.getcwd(), "chromedriver", f"_{chrome_version}")
        chromedriver_path = download_chromedriver(chrome_version, download_dir)

        # 5. 验证版本匹配
        success, message = verify_chromedriver(chrome_version, chromedriver_path)
        print(message)

        if success:
            print("\n操作成功完成！ChromeDriver 已准备就绪。")
        else:
            print("\n警告: ChromeDriver 版本可能不兼容，请手动验证。")

        return 0 if success else 1
    except Exception as e:
        print(f"\n错误: {str(e)}")
        return 1


if __name__ == "__main__":
    exit_code = main()
    # 防止窗口立即关闭（如果双击运行）
    if exit_code != 0:
        input("\n按 Enter 键退出...")
